import ee
import os
from ee.batch import Export, Task
import time

from common.paii_gee import gee_utils, gcp_utils
from common.logger import logger

os.environ[
    "GOOGLE_APPLICATION_CREDENTIALS"
] = "/NAS6/Members/linchenxi/Earth Engine default project-d372e5cfe621.json"

try:
    ee.Initialize()
except Exception as e:
    ee.Authenticate()
    ee.Initialize()


def pipeline(mode="tile"):
    if "0" in ops_idx:
        logger.info(f"#----- Step 0. {operations['0']}-----#")
        gcp_utils.delete_blob(bucket_name)

    if "1" in ops_idx:
        logger.info(f"#----- Step 1. {operations['1']}-----#")
        for img in os.listdir(server_folder_upload):
            if mode == "tile" and img.split(".")[0].endswith("class"):
                gcp_utils.upload_blob(
                    "paii_gee", os.path.join(server_folder_upload, img), f"upload/{img}"
                )
            elif mode == "tile" and not img.split(".")[0].endswith("class"):
                logger.info(
                    f"File {os.path.join(server_folder_upload, img)} is being skipped."
                )
            else:
                gcp_utils.upload_blob(
                    "paii_gee", os.path.join(server_folder_upload, img), f"upload/{img}"
                )

    if "2" in ops_idx:
        logger.info(f"#----- Step 2. {operations['2']}-----#")
        if mode == "tile":
            file_num = len(os.listdir(server_folder_upload))
        else:
            file_num = len(os.listdir(server_folder_upload))
        while len(gcp_utils.list_blobs(bucket_name)) != file_num:
            logger.info("Files are still uploading, please wait.")
        # when all files have been uploaded tp Google Cloud Platform
        file_to_upload = gcp_utils.list_blobs(bucket_name)
        for idx, item in enumerate(file_to_upload):
            if "upload" in item and item.endswith(".tif"):
                if item.split("/")[-1].split(".")[0].split("_")[-1] == "class":
                    logger.info(
                        f"The naming follows the tile mode for {item.split('/')[-1]}"
                    )
                    mode = "tile"
                    tile = item.split("/")[-1].split(".")[0].split("_")[0]
                    year = item.split("/")[-1].split(".")[0].split("_")[1]
                    identifier = item.split("/")[-1].split(".")[0].split("_")[2]
                    properties = {"Tile": tile, "Year": year, "Identifier": identifier}
                else:
                    logger.info(
                        f"The naming follows the city mode for {item.split('/')[-1]}"
                    )
                    mode = "city"
                    city = item.split("/")[-1].split(".")[0].split("_")[-2]
                    county = item.split("/")[-1].split(".")[0].split("_")[-1]
                    province = item.split("/")[-1].split(".")[0].split("_")[1]
                    properties = {"County": county, "City": city, "Province": province}
                if not gee_utils.AssetManager.asset_exists(gee_asset_name):
                    ee.data.createAsset({"type": "ImageCollection"}, gee_asset_name)
                gcp_utils.gee_upload_image(
                    os.path.join(gee_asset_name, str(idx)),
                    f"gs://paii_gee/" + item,
                    properties,
                )

    if "3" in ops_idx:
        logger.info(f"#----- Step 3. {operations['3']}-----#")
        active_task_list = gee_utils.TaskManager.get_active_task()
        active_task_list = [task.state for task in active_task_list]
        while len(active_task_list) != 0:
            logger.info(f"{len(active_task_list)} are still runing. Please wait.")
            time.sleep(30)
            active_task_list = gee_utils.TaskManager.get_active_task()
            active_task_list = [task.state for task in active_task_list]
        crop_layer = ee.ImageCollection(gee_asset_name)
        image_smoother = gee_utils.ImageSmoother(
            imgCol=crop_layer,
            areaThreshold=1000,
            num_cls=len(target_crop),
            circleRadius=10,
        )
        pp_data = image_smoother.deSpeckler().map(
            lambda img: img.reduceNeighborhood(
                reducer=ee.Reducer.mode(),
                kernel=ee.Kernel.square(12.5, "meters"),
                optimization="window",
            )
        )

    if "4" in ops_idx:
        logger.info(f"#----- Step 4. {operations['4']}-----#")
        pp_data_lst = pp_data.toList(1000)
        ref_lst = crop_layer.toList(1000)
        for i in range(pp_data.size().getInfo()):
            export_img = ee.Image(pp_data_lst.get(i))
            ref_img = ee.Image(ref_lst.get(i))
            if mode == "tile":
                filename = (
                    ee.String(ref_img.get("Tile")).getInfo()
                    + "_"
                    + ee.String(ref_img.get("Year")).getInfo()
                    + "_"
                    + ee.String(ref_img.get("Identifier")).getInfo()
                )
            if mode == "city":
                if (
                    ee.String(ref_img.get("Province")).getInfo()
                    != ee.String(ref_img.get("City")).getInfo()
                ):
                    filename = (
                        "中华人民共和国_"
                        + ee.String(ref_img.get("Province")).getInfo()
                        + "_"
                        + ee.String(ref_img.get("City")).getInfo()
                        + "_"
                        + ee.String(ref_img.get("County")).getInfo()
                    )
                else:
                    filename = (
                        "中华人民共和国_"
                        + ee.String(ref_img.get("Province")).getInfo()
                        + "_"
                        + ee.String(ref_img.get("County")).getInfo()
                    )
            task = Export.image.toCloudStorage(
                image=export_img,
                description=f"exporting {i}th result",
                bucket="paii_gee",
                fileNamePrefix=os.path.join("download", filename),
                region=export_img.geometry(),
                maxPixels=1e13,
                scale=10,
            )
            task.start()

    if "5" in ops_idx:
        logger.info(f"#----- Step 5. {operations['5']}-----#")
        active_task_list = gee_utils.TaskManager.get_active_task()
        active_task_list = [task.state for task in active_task_list]
        while len(active_task_list) != 0:
            logger.info(f"{len(active_task_list)} are still runing. Please wait.")
            time.sleep(300)
            active_task_list = gee_utils.TaskManager.get_active_task()
            active_task_list = [task.state for task in active_task_list]
        gcp_utils.download_blob(
            "paii_gee",
            server_folder_download,
        )

    if "6" in ops_idx:
        from common.geoimage.raster_dataset import mosaic_raster_files
        from catalog.place import PlaceAPI

        city_count = {}
        for item in os.listdir(server_folder_download):
            city = item.split(".")[0]
            for idx, s in enumerate(city):
                if s == "0":
                    idx -= 1
                    break
            city = city[: idx + 1]
            if city not in city_count.keys():
                city_count[city] = 1
            else:
                city_count[city] += 1
        for city in city_count.keys():
            if city_count[city] == 1:
                continue
            image_lst = [
                os.path.join(server_folder_download, item)
                for item in os.listdir(server_folder_download)
                if city in item
            ]
            # geom = PlaceRecord.query_many_items(target_parent_name=city.split("_")[2], region_level=2)
            geom = PlaceAPI.get_geom_list_with_place_name(
                target_name=city.split("_")[2], target_parent_name=city.split("_")[1]
            )[0]
            mosaic = mosaic_raster_files(
                image_lst,
                dest_aoi=geom,
                crop_to_aoi=True,
                fpath_dest=os.path.join(server_folder_download, city + ".tif"),
            )


if __name__ == "__main__":
    operations = {
        "0": "delete all files in gcp",
        "1": "upload from server to gcp",
        "2": "upload from gcp to gee",
        "3": "post-process on gee",
        "4": "download from gee to gcp",
        "5": "download from gcp to server",
        "6": "mosaic large files",
    }
    ops_idx = ["0", "1", "2"]
    bucket_name = "paii_gee"
    gee_asset_name = "users/lin00370/paii/xinjiang/cotton_model1_filtered"
    server_folder_upload = "/NAS6/Members/linchenxi/projects/crop_recognition/inference/xinjiang/model_1/filtered_city/2023"
    server_folder_download = "/NAS6/Members/linchenxi/projects/crop_recognition/inference/liaoning/model_2/filtered_city/2023"
    areaThreshold = 1000
    circleRadius = 10
    crop_of_interest = {
        "heilongjiang": ["negative", "corn", "soybeans", "rice"],
        "liaoning": ["negative", "corn", "soybeans", "rice"],
        "hubei": ["negative", "rice"],
        "innermongolia": ["negative", "corn", "soybeans"],
        "xinjiang": ["negative", "cotton"],
        "henan": ["negative", "corn"],
        "shandong": ["negative", "corn"],
    }
    target_crop = crop_of_interest["liaoning"]
    mode = "city"
    pipeline(mode)
